package main

import (
	"cs466/findmotif/fasta"
	"cs466/findmotif/gen"
	"flag"
	"fmt"
	"log"
	"os"
	"encoding/json"
)

var out = flag.String("out", "./out", "Output directory")
var tasks_json = flag.String("tasks", "./default_tasks.json", "Task specification file")
var n = flag.Int("n", 10, "Datasets to generate per task")

var defaults = gen.Task{
	ICPC:        2,
	MotifLength: 8,
	SeqLength:   500,
	SeqCount:    10,
}

func WriteDataset(t gen.Task, n int) error {
	seqs, sites, pwm, err := t.Generate()
	if err != nil {
		panic(err) // err can only be I/O errors on /dev/urandom, which should never happen
	}

	// make the dataset directory
	dirname := fmt.Sprintf("%s/%.1f-%d-%d-%d-%d", *out, t.ICPC, t.MotifLength, t.SeqLength, t.SeqCount, n)
	err = os.MkdirAll(dirname, os.ModeDir|0755)
	if err != nil {
		return err
	}

	// write the FASTA sequence file
	fasta_name := fmt.Sprintf("%s/sequences.fa", dirname)
	f, err := os.Create(fasta_name)
	if err != nil {
		return err
	}
	defer f.Close()
	err = fasta.WriteFastaSlice(f, seqs)
	if err != nil {
		return err
	}

	// record the location of the planted motifs
	sites_name := fmt.Sprintf("%s/sites.txt", dirname)
	f, err = os.Create(sites_name)
	if err != nil {
		return err
	}
	defer f.Close()
	_, err = fmt.Fprint(f, sites)
	if err != nil {
		return err
	}

	// record the motif matrix
	motif_name := fmt.Sprintf("%s/motif.txt", dirname)
	f, err = os.Create(motif_name)
	if err != nil {
		return err
	}
	defer f.Close()
	err = fasta.WriteMotif(f, pwm, t.MotifLength, float64(t.SeqCount))
	if err != nil {
		return err
	}

	// record the motif length
	ml_name := fmt.Sprintf("%s/motiflength.txt", dirname)
	f, err = os.Create(ml_name)
	if err != nil {
		return err
	}
	defer f.Close()
	_, err = fmt.Fprint(f, t.MotifLength)
	if err != nil {
		return err
	}

	return nil
}

func DoDatasets(t gen.Task, n int) {
	for i := 0; i < n; i++ {
		if err := WriteDataset(t, i); err != nil {
			log.Fatal(err)
		}
	}
}

func main() {
	flag.Parse()
	_, err := os.Stat(*out)
	if err != nil && !os.IsNotExist(err) {
		log.Fatal(err)
	}
	if !os.IsNotExist(err) {
		log.Fatalf("Output directory %q already exists", *out)
	}
	os.MkdirAll(*out, os.ModeDir|0755)

	f, err := os.Open(*tasks_json)
	if err != nil {
		log.Fatal(err)
	}

	var tasks []gen.Task
	json.NewDecoder(f).Decode(&tasks)

	for _, t := range tasks {
		DoDatasets(t, *n)
	}
}
